from __future__ import annotations

import torch
import torch.nn as nn


class MnistLeNet(nn.Module):
    """A small CNN suitable for MNIST (LeNet-style).

    Input: (N, 1, 28, 28)
    Output: (N, 10)
    
    Args:
        num_classes (int): Number of output classes. Defaults to 10.
    """

    def __init__(self, num_classes: int = 10) -> None:
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(inplace=True),
            nn.Conv2d(32, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(inplace=True),
            nn.MaxPool2d(2),  # 14x14
            nn.Dropout(0.25),

            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(inplace=True),
            nn.MaxPool2d(2),  # 7x7
            nn.Dropout(0.25),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),  # 64 * 7 * 7 = 3136
            nn.Linear(64 * 7 * 7, 256),
            nn.LeakyReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(256, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        x = self.classifier(x)
        return x


def build_model(num_classes: int = 10) -> MnistLeNet:
    """Build and return a new MnistLeNet model instance.

    Args:
        num_classes (int): Number of output classes. Defaults to 10.

    Returns:
        MnistLeNet: A new model instance.
    """
    return MnistLeNet(num_classes=num_classes)


def load_model(weights_path: str, device: torch.device | str = "cpu", num_classes: int = 10) -> MnistLeNet:
    """Load a pre-trained MnistLeNet model from a checkpoint file.

    Args:
        weights_path (str): Path to the checkpoint file.
        device (Union[torch.device, str]): Device to load the model on. Defaults to "cpu".
        num_classes (int): Number of output classes. Defaults to 10.

    Returns:
        MnistLeNet: The loaded model instance.

    Raises:
        FileNotFoundError: If the checkpoint file does not exist.
        RuntimeError: If the model loading fails.
    """
    try:
        device = torch.device(device)
        model = build_model(num_classes=num_classes)
        checkpoint = torch.load(weights_path, map_location=device)
        if isinstance(checkpoint, dict) and "model_state" in checkpoint:
            model.load_state_dict(checkpoint["model_state"])
        else:
            model.load_state_dict(checkpoint)
        model.to(device)
        model.eval()
        return model
    except FileNotFoundError:
        raise FileNotFoundError(f"Checkpoint file not found: {weights_path}")
    except Exception as e:
        raise RuntimeError(f"Failed to load model: {str(e)}")